import os
import sys

from PyQt5.QtGui import QPixmap
from PyQt5.QtWidgets import QApplication, QMainWindow, QMessageBox, QFileDialog, QWidget, QVBoxLayout, QLabel

from ui.ui_main import Ui_Main
from worker.TestWorker import TestWorker
from worker.TrainWorker import TrainWorker
from worker.ViewWorker import ViewWorker


class ImageWindow(QWidget):
    def __init__(self, image_path):
        super().__init__()
        self.setWindowTitle("图片展示窗口")
        self.resize(400, 300)

        layout = QVBoxLayout()

        self.image_label = QLabel()
        pixmap = QPixmap(image_path)
        self.image_label.setPixmap(pixmap)
        self.image_label.setScaledContents(True)  # 图片自适应标签大小

        layout.addWidget(self.image_label)
        self.setLayout(layout)


class MainApp(QMainWindow, Ui_Main):
    def __init__(self):
        super().__init__()
        self.image_window = None
        self.setupUi(self)  # 加载界面
        self.testWorker = None
        self.trainWorker = None
        self.training.clicked.connect(lambda: self.stackedWidget.setCurrentWidget(self.page0))  # 打开模型训练窗口
        self.content.clicked.connect(lambda: self.stackedWidget.setCurrentWidget(self.page1))  # 打开内容分发窗口
        # 开始训练模型
        self.begin.clicked.connect(self.on_begin_click)
        # 查看模型训练过程
        self.search_log.clicked.connect(self.on_search_click)

        # 查看模型可视化界面
        self.pushButton_3.clicked.connect(self.on_pushButton_3_click)

        # 开始内容转发过程
        self.pushButton.clicked.connect(self.on_forward_click)

        # 查看日志
        self.pushButton_2.clicked.connect(self.on_search_click)

        # 查看网络拓扑
        self.pushButton_4.clicked.connect(self.open_image_window)

    def on_begin_click(self):
        """开始训练模型"""

        # 1. 获取输入
        combo_text = self.comboBox.currentText().strip()
        combo_text2 = self.comboBox_2.currentText().strip()
        combo_text3 = self.comboBox_3.currentText().strip()
        line_text = self.lineEdit.text().strip()
        line2_text = self.lineEdit_2.text().strip()
        line3_text = self.lineEdit_3.text().strip()

        if not combo_text or not combo_text2 or not combo_text3:
            QMessageBox.warning(self, "提示", "下拉框内容不能为空！")
            return
        if not line_text or not line2_text or not line3_text:
            QMessageBox.warning(self, "提示", "文本框内容不能为空！")
            return

        # 清空上次输出
        self.textBrowser.clear()

        params = {
            "script": "train.py",
            "config_file": combo_text,
            "datadir": combo_text2,
            "dataset_config_file": combo_text3,
            "pp": line_text,
            "train_batch_size": line2_text,
            "max_epochs": line3_text
        }
        self.trainWorker = TrainWorker(params)
        self.trainWorker.output_signal.connect(lambda text: self.textBrowser.append(text))
        self.trainWorker.progress_signal.connect(lambda progress: self.progressBar.setValue(progress))
        self.trainWorker.start()

    def on_search_click(self):
        # 获取程序当前工作目录下的某个相对路径
        base_dir = os.getcwd()  # 或者 os.path.dirname(__file__)
        relative_path = os.path.join(base_dir, "output")
        # 弹出文件选择对话框，起始路径设为相对目录
        file_path, _ = QFileDialog.getOpenFileName(
            self,
            "选择文件",
            relative_path,
            "所有文件 (*.*)"  # 可以根据需要设置过滤器，如 "图像文件 (*.png *.jpg *.jpeg)"
        )
        if file_path:
            os.startfile(file_path)
            # self.textBrowser.append(file_path)

    def on_pushButton_3_click(self):
        self.viewWorker = ViewWorker()
        self.viewWorker.output_signal.connect(lambda text: self.textBrowser.append(text))
        self.viewWorker.start()

    def on_search2_click(self):
        file_path, _ = QFileDialog.getOpenFileName(self, '选择文件', '', '所有文件 (*);;文本文件 (*.txt)')
        if file_path:
            self.textBrowser_2.append(file_path)

    def on_forward_click(self):
        combo_text4 = self.comboBox_4.currentText().strip()
        combo_text5 = self.comboBox_5.currentText().strip()
        combo_text6 = self.comboBox_6.currentText().strip()
        combo_text7 = self.comboBox_7.currentText().strip()
        combo_text8 = self.comboBox_8.currentText().strip()
        combo_text9 = self.comboBox_9.currentText().strip()
        line5_text = self.lineEdit_5.text().strip()

        if (not combo_text4 or not combo_text5 or not combo_text6 or not combo_text7 or not combo_text8
                or not combo_text9):
            QMessageBox.warning(self, "提示", "下拉框内容不能为空！")
            return
        if not line5_text:
            QMessageBox.warning(self, "提示", "文本框内容不能为空！")
            return

        # 清空上次输出
        self.textBrowser_2.clear()

        params = {
            "script": "my_val.py",
            "config_file": combo_text7,
            "datadir": combo_text8,
            "dataset_config_file": combo_text9,
            "pretrained": combo_text4,
            "net_path": combo_text5,
            "print_freq": line5_text
        }
        # 启动子线程
        self.testWorker = TestWorker(params, no_cache=(combo_text6 == '否'))
        # 信号连接
        self.testWorker.output_signal.connect(lambda text: self.textBrowser_2.append(text))
        self.testWorker.progress_signal.connect(lambda value: self.progressBar_2.setValue(value))
        self.testWorker.start()

    def open_image_window(self):
        image_path = "docs/内容转发网络拓扑.jpg"
        self.image_window = ImageWindow(image_path)
        self.image_window.show()


if __name__ == '__main__':
    app = QApplication(sys.argv)
    myWin = MainApp()
    myWin.show()
    sys.exit(app.exec_())
